{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "99cea58c-48bc-4af6-8358-df9695659983",
            "metadata": {
                "tags": []
            },
            "source": [
                "# Retrieval-Augmented OpenAI Agent"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "673df1fe-eb6c-46ea-9a73-a96e7ae7942e",
            "metadata": {
                "tags": []
            },
            "source": [
                "In this tutorial, we show you how to use our `FnRetrieverOpenAI` implementation\n",
                "to build an agent on top of OpenAI's function API and store/index an arbitrary number of tools. Our indexing/retrieval modules help to remove the complexity of having too many functions to fit in the prompt."
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "54b7bc2e-606f-411a-9490-fcfab9236dfc",
            "metadata": {
                "tags": []
            },
            "source": [
                "## Initial Setup "
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "23e80e5b-aaee-4f23-b338-7ae62b08141f",
            "metadata": {},
            "source": [
                "Let's start by importing some simple building blocks.  \n",
                "\n",
                "The main thing we need is:\n",
                "1. the OpenAI API\n",
                "2. a place to keep conversation history \n",
                "3. a definition for tools that our agent can use."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "id": "9d47283b-025e-4874-88ed-76245b22f82e",
            "metadata": {
                "tags": []
            },
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "/Users/suo/miniconda3/envs/llama/lib/python3.9/site-packages/deeplake/util/check_latest_version.py:32: UserWarning: A newer version of deeplake (3.6.7) is available. It's recommended that you update to the latest version using `pip install -U deeplake`.\n",
                        "  warnings.warn(\n"
                    ]
                }
            ],
            "source": [
                "import json\n",
                "from typing import Sequence\n",
                "\n",
                "from llama_index.tools import BaseTool, FunctionTool"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "6fe08eb1-e638-4c00-9103-5c305bfacccf",
            "metadata": {},
            "source": [
                "Let's define some very simple calculator tools for our agent."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "id": "3dd3c4a6-f3e0-46f9-ad3b-7ba57d1bc992",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "def multiply(a: int, b: int) -> int:\n",
                "    \"\"\"Multiply two integers and returns the result integer\"\"\"\n",
                "    return a * b\n",
                "\n",
                "\n",
                "def add(a: int, b: int) -> int:\n",
                "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
                "    return a + b\n",
                "\n",
                "\n",
                "def useless(a: int, b: int) -> int:\n",
                "    \"\"\"Toy useless function.\"\"\"\n",
                "    pass\n",
                "\n",
                "\n",
                "multiply_tool = FunctionTool.from_defaults(fn=multiply, name=\"multiply\")\n",
                "useless_tools = [\n",
                "    FunctionTool.from_defaults(fn=useless, name=f\"useless_{str(idx)}\")\n",
                "    for idx in range(28)\n",
                "]\n",
                "add_tool = FunctionTool.from_defaults(fn=add, name=\"add\")\n",
                "\n",
                "all_tools = [multiply_tool] + [add_tool] + useless_tools\n",
                "all_tools_map = {t.metadata.name: t for t in all_tools}"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "8170dd32-fa00-458c-aeb5-37292d0773c0",
            "metadata": {},
            "source": [
                "## Building an Object Index\n",
                "\n",
                "We have an `ObjectIndex` construct in LlamaIndex that allows the user to use our index data structures over arbitrary objects.\n",
                "The ObjectIndex will handle serialiation to/from the object, and use an underying index (e.g. VectorStoreIndex, SummaryIndex, KeywordTableIndex) as the storage mechanism. \n",
                "\n",
                "In this case, we have a large collection of Tool objects, and we'd want to define an ObjectIndex over these Tools.\n",
                "\n",
                "The index comes bundled with a retrieval mechanism, an `ObjectRetriever`. \n",
                "\n",
                "This can be passed in to our agent so that it can \n",
                "perform Tool retrieval during query-time."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "id": "6704a755-7f05-43a3-8a56-f5f587ae4c40",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "# define an \"object\" index over these tools\n",
                "from llama_index import VectorStoreIndex\n",
                "from llama_index.objects import ObjectIndex, SimpleToolNodeMapping\n",
                "\n",
                "tool_mapping = SimpleToolNodeMapping.from_objects(all_tools)\n",
                "obj_index = ObjectIndex.from_objects(\n",
                "    all_tools,\n",
                "    tool_mapping,\n",
                "    VectorStoreIndex,\n",
                ")"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "707d30b8-6405-4187-a9ed-6146dcc42167",
            "metadata": {
                "tags": []
            },
            "source": [
                "## Our `FnRetrieverOpenAIAgent` Implementation "
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "798ca3fd-6711-4c0c-a853-d868dd14b484",
            "metadata": {},
            "source": [
                "We provide a `FnRetrieverOpenAIAgent` implementation in LlamaIndex, which can take in an `ObjectRetriever` over a set of `BaseTool` objects.\n",
                "\n",
                "During query-time, we would first use the `ObjectRetriever` to retrieve a set of relevant Tools. These tools would then be passed into the agent; more specifically, their function signatures would be passed into the OpenAI Function calling API. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "id": "38ab3938-1138-43ea-b085-f430b42f5377",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "from llama_index.agent import FnRetrieverOpenAIAgent"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "id": "d852ece7-e5a1-4368-9d59-c7014e0b5b4d",
            "metadata": {
                "tags": []
            },
            "outputs": [],
            "source": [
                "agent = FnRetrieverOpenAIAgent.from_retriever(obj_index.as_retriever(), verbose=True)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 11,
            "id": "33ea069f-819b-4ec1-a93c-fcbaacb362a1",
            "metadata": {
                "tags": []
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "=== Calling Function ===\n",
                        "Calling function: multiply with args: {\n",
                        "  \"a\": 212,\n",
                        "  \"b\": 122\n",
                        "}\n",
                        "Got output: 25864\n",
                        "========================\n"
                    ]
                },
                {
                    "data": {
                        "text/plain": [
                            "Response(response='212 multiplied by 122 is 25,864.', source_nodes=[], metadata=None)"
                        ]
                    },
                    "execution_count": 11,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "agent.chat(\"What's 212 multiplied by 122? Make sure to use Tools\")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 12,
            "id": "ec423b90-59cd-40ef-b497-a3842b3e7b58",
            "metadata": {
                "tags": []
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "=== Calling Function ===\n",
                        "Calling function: add with args: {\n",
                        "  \"a\": 212,\n",
                        "  \"b\": 122\n",
                        "}\n",
                        "Got output: 334\n",
                        "========================\n"
                    ]
                },
                {
                    "data": {
                        "text/plain": [
                            "Response(response='212 added to 122 is 334.', source_nodes=[], metadata=None)"
                        ]
                    },
                    "execution_count": 12,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "agent.chat(\"What's 212 added to 122 ? Make sure to use Tools\")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "65bf017c-d381-414c-a4d4-d6a9d848e186",
            "metadata": {},
            "outputs": [],
            "source": []
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3 (ipykernel)",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.9.16"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}